package com.bawei.day05;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.operator.common.evaluation.TuningMultiClassMetric;
import com.alibaba.alink.pipeline.classification.RandomForestClassifier;
import com.alibaba.alink.pipeline.tuning.*;
import org.apache.flink.types.Row;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class GridSearchCVTest1 {
	@Test
	public void testRandomForestClassifier() throws Exception {
		//1、
		List <Row> df = Arrays.asList(
				Row.of(1.0, "A", 0, 0, 0),
				Row.of(2.0, "B", 1, 1, 0),
				Row.of(3.0, "C", 2, 2, 1),
				Row.of(4.0, "D", 3, 3, 1)
		);

		BatchOperator <?> batchSource
				= new MemSourceBatchOp(df, " f0 double, f1 string, f2 int, f3 int, label int");
		//2、
		RandomForestClassifier rf = new RandomForestClassifier()
				.setPredictionDetailCol("pred_detail")
				.setPredictionCol("pred")
				.setLabelCol("label")
				.setFeatureCols("f0", "f1", "f2", "f3");
		//3、调参
		ParamGrid paramGrid = new ParamGrid()
				.addGrid(rf, RandomForestClassifier.SUBSAMPLING_RATIO, new Double[] {1.0, 0.99})
				.addGrid(rf, RandomForestClassifier.NUM_TREES, new Integer[] {3, 4});


		// 二分类评估报错
		/*BinaryClassificationTuningEvaluator tuningEvaluator = new BinaryClassificationTuningEvaluator()
                .setLabelCol("label")
                .setPredictionDetailCol("pred_detail")
                .setTuningBinaryClassMetric(TuningBinaryClassMetric.ACCURACY);*/
		// 多分类评估
		MultiClassClassificationTuningEvaluator tuningEvaluator = new MultiClassClassificationTuningEvaluator()
				.setLabelCol("label")
				.setPredictionDetailCol("pred_detail")
				.setTuningMultiClassMetric(TuningMultiClassMetric.ACCURACY);

		// 网络搜索
		GridSearchCV cv = new GridSearchCV()
				.setEstimator(rf)
				.setParamGrid(paramGrid)
				.setTuningEvaluator(tuningEvaluator)
				.setNumFolds(2)
				.enableLazyPrintTrainInfo("TrainInfo");
		GridSearchCVModel model = cv.fit(batchSource);

		//提取最佳模型预测
		BatchOperator<?> result = model.getBestPipelineModel().transform(batchSource);
		result.print();


	}
}
